import torch
import torch.nn as nn
from perceptron_model import PerceptronModel
from net_model import MultiLayerPerceptronModel

# 数据
X = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0], dtype=torch.float32)
y = torch.tensor([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0], dtype=torch.float32)

# 设置GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:{}".format(device))
# 单层感知机模型训练
model1 = PerceptronModel(1, 1).to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model1.parameters(), lr=0.01)

num_epochs = 500
for epoch in range(num_epochs):
    inputs = X.unsqueeze(1).to(device)
    targets = y.unsqueeze(1).to(device)
    
    outputs = model1(inputs)
    loss = criterion(outputs, targets)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
        
# 保存单层感知机模型参数
torch.save(model1.state_dict(), 'perceptron_model.pth')
print("Saved perceptron model to perceptron_model.pth")
# 多层感知机模型训练
model2 = MultiLayerPerceptronModel(1, 2, 1).to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model2.parameters(), lr=0.01)

num_epochs = 500
for epoch in range(num_epochs):
    inputs = X.unsqueeze(1).to(device)
    targets = y.unsqueeze(1).to(device)
    
    outputs = model2(inputs)
    loss = criterion(outputs, targets)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
        
# 保存多层感知机模型参数        
torch.save(model2.state_dict(), 'net_model.pth')
print("Saved net model to net_model.pth")